# will you just take some input and ask for EOF?
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-n","--name",type=str,required=True)
parser.add_argument("-d","--decode",action="store_true")
parser.add_argument("-p","--password",type=str,required=True)

args = parser.parse_args()
import os
import uuid
import json
import random
import base64

def encode(string):
    if type(string) == str:
        string = string.encode("utf-8")
    return base64.b64encode(string).decode("utf-8")
def decode(string):
    if type(string) == str:
        string = string.encode("utf-8")
    return base64.b64decode(string)


def get_password(password):
    assert type(password) == str
    l = len(password)
    #print(l)
    assert l > 0 and l <=16
    add = 16 - l
    return password + " "*add

def get_secret():
    if os.path.exists(".secret"):
        secret = open(".secret","r").read()
        assert len(secret) == 16
        return secret
    else:
        uuid4 = str(uuid.uuid4())
        uuid4 = uuid4.replace("-","")
        uuid4 = [x for x in uuid4]
        random.shuffle(uuid4)
        uuid4 = uuid4[:16]
        uuid4 = "".join(uuid4)
        with open(".secret","w+") as f:
            f.write(uuid4)
        return uuid4

"""
ECB没有偏移量
"""
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex


def add_to_16(text):
    #print(type(text))
    if len(text.encode('utf-8')) % 16:
        add = 16 - (len(text.encode('utf-8')) % 16)
    else:
        add = 0
    text = text + ('\0' * add)
    return text.encode('utf-8')


# 加密函数
def encrypt(text,password):
    password = get_password(password)
    key = password.encode('utf-8')
    mode = AES.MODE_ECB
    text = add_to_16(text)
    cryptos = AES.new(key, mode)

    cipher_text = cryptos.encrypt(text)
    return b2a_hex(cipher_text)


# 解密后，去掉补足的空格用strip() 去掉
def decrypt(text,password):
    password = get_password(password)
    key = password.encode('utf-8')
    mode = AES.MODE_ECB
    cryptor = AES.new(key, mode)
    plain_text = cryptor.decrypt(a2b_hex(text))
    return bytes.decode(plain_text).rstrip('\0')


assert args.name is not None
assert args.name not in [".secret","vault.py","vault.sh","vault_bin.py", "vault_bin.sh","test.sh","test.log"]
assert os.path.exists(args.name)
def verify(content):
    content = json.loads(content)
    sign = content["sign"]
    ctype = content["type"]
    assert ctype == "text"
    sign = decode(sign)
    assert list(content.keys()) == ["content","sign","challenge","type"]
    content = content["content"]
    content = decode(content)
    message = decrypt(sign,get_secret())
    message = decode(message)
    assert message == content

if not args.decode:
    with open(args.name, "r") as f:
        content = f.read()
#        content0 = content
        try:
            verify(content)
            print("Already encrypted. Cannot proceed.")
        except:
            #print(content,type(content))
            challenge = str(uuid.uuid4())
            content = {"content":content,"challenge":challenge}
            content = json.dumps(content)
            content = encode(content)
            content = encrypt(content,args.password)
            secret = get_secret()
            content = encode(content)
            sign = encrypt(content,secret)
            sign = encode(sign)
            content = {"content":content,"sign":sign,"challenge":challenge,"type":"text"}
            content = json.dumps(content)
            with open(args.name,"w+") as f:
                f.write(content)
else:
    with open(args.name, "r") as f:
        content = f.read()
        verify(content)
        content = json.loads(content)
        target = content["challenge"]
        content = content["content"]
        content = decode(content)
        #print(content)
        content = decrypt(content, args.password)
        content = decode(content)
        #print(content)
        content = json.loads(content)
        assert list(content.keys()) == ["content","challenge"]
        challenge = content["challenge"]
        content = content["content"]
        assert target == challenge
        print("___________DECRYPTED MESSAGE___________")
        print(content, end="")
        # prevent double encryption?
    #print("content:")
    #print(content)

